import numpy as np
import cv2
from scipy.misc import imresize
from moviepy.editor import VideoFileClip
# from IPython.display import HTML
from keras.models import load_model

# Load Keras model
model = load_model('model/lane_model_eps10.h5', compile=False)
# model = load_model('model/full_CNN_model_impromenteps15.h5', compile=False)

# Class to average lanes with
class Lanes():
    def __init__(self):
        self.recent_fit = []
        self.avg_fit = []


def road_lines(image):
    """ Takes in a road image, re-sizes for the model,
    predicts the lane to be drawn from the model in G color,
    recreates an RGB image of a lane and merges with the
    original road image.
    """

    # Get image ready for feeding into model
    small_img = imresize(image, (80, 160, 3))
    # small_img = np.array(Image.fromarray(image).resize((3,(80, 160))))
    small_img = np.array(small_img)
    small_img = small_img[None, :, :, :]

    # Make prediction with neural network (un-normalize value by multiplying by 255)
    prediction = model.predict(small_img)[0] * 255
    # print(prediction)

    # Add lane prediction to list for averaging
    lanes.recent_fit.append(prediction)
    # Only using last five for average
    if len(lanes.recent_fit) > 5:
        lanes.recent_fit = lanes.recent_fit[1:]

    # Calculate average detection
    lanes.avg_fit = np.mean(np.array([i for i in lanes.recent_fit]), axis=0)

    # Generate fake R & B color dimensions, stack with G
    blanks = np.zeros_like(lanes.avg_fit).astype(np.uint8)
    lane_drawn = np.dstack((blanks, lanes.avg_fit, blanks))

    # Re-size to match the original image
    lane_image = imresize(lane_drawn, (720, 1280, 3))
    # lane_image = np.array(Image.fromarray(lane_drawn).resize((3,(1080,1920))))
    # Merge the lane drawing onto the original image
    result = cv2.addWeighted(image, 1, lane_image, 1, 0)

    # cv2.imshow('d', lanes.avg_fit)
    # cv2.waitKey(0)
    return result, lane_image


lanes = Lanes()


def detect_image(image):
    # image = cv2.imread(image_path)
    image = cv2.resize(image, (1280, 720), interpolation=cv2.INTER_LINEAR)
    result, lane_image = road_lines(image)
    return result, lane_image, image

if __name__ == '__main__':
    image = cv2.imread('new_img.jpg')
    res = detect_image(image)
